import os
import random
import time

from dataset import load_dataset
from lazyselect import lazyselect
from quickselect import quickselect
from utils import brief, save_variable

if __name__ == '__main__':
    run_num = 1000
    name = 'lazyselect'

    # select = quickselect
    select = lazyselect

    print('name:', name)
    print('select:', select)

    dataset_sizes = [1000, 2000, 5000, 10000, 20000, 50000, 100000]
    avg_consume = []
    rand = random.Random()
    rand.seed(5)

    for size in dataset_sizes:
        s = load_dataset(os.path.join('datasets', f'dataset-{str(size)}.txt'))
        print('dataset size:', len(s))

        k_list = rand.choices(range(1, size + 1), k=run_num)
        # print('k_list:', brief(k_list), ', len =', len(k_list))

        # s_ordered = sorted(s)
        # correct_value_list = [s_ordered[k - 1] for k in k_list]
        # print('correct_value_list:', brief(correct_value_list, 2), ', len =', len(correct_value_list))

        start = time.time()
        # for k, correct_value in zip(k_list, correct_value_list):
        for k in k_list:
            result = select(s, k)
            # if result != correct_value:
            #     print("Wrong result!!!")
        elapsed = time.time() - start
        elapsed = elapsed / run_num
        avg_consume.append(elapsed)
        print("Average time consume:", elapsed * 1000, 'ms')
        print('=' * 50)

    output = {
        'dataset_sizes': dataset_sizes,
        'avg_time_consume': avg_consume
    }
    os.makedirs('results', exist_ok=True)
    save_variable(output, os.path.join('results', name + '.pkl'))

    print("Finished.")
